import logging
import shutil
from datetime import datetime

from .catboost_model.predict_data import *
from .catboost_model.training_data import train_main as nonlinear_train

from .linear_model import train as linear_train
from .linear_model import predict_strength, config
from django.http import JsonResponse
from functools import wraps
from .platform_api_handler import PlatformAPIHandler
from django.conf import settings

import sys

sys.path.append(settings.BASE_DIR)
from celery.result import AsyncResult
from redis_django_celery.base_info.task import linear_train_async, nonlinear_train_async  # TODO:使用异步通讯


def logging_decorator(func):
    """
    打印输出日志修饰器
    :param func:
    :return:
    """

    @wraps(func)
    def wrapper(request):
        current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        func_name = func.__name__
        logging.info(f"[{current_time}] Function: {func_name} - Path: {request.META['PATH_INFO']}")
        json_content = func(request)
        logging.info(f"[{current_time}] Function: {func_name} - Response: {json_content}")
        return JsonResponse(json_content, safe=False)

    return wrapper


# ========================  需要报漏出去的两个API接口 =====================
@logging_decorator
def cement_predict(request):
    """
    模型预测接口
    :param request:
    :return:
    """
    # TODO 需要带上模型的版本号??

    if request.method == "POST":
        results = []
        # 获取请求的原始 body 数据
        body_data = request.body.decode('utf-8')
        json_data = json.loads(body_data)
        data = json_data.get("data", {})
        model_code = f"CI_10001194-{data['cement_mill']}-{data['cement_type']}"
        data = predict_strength.transform_ci_to_nl(data, config.master_data)
        # 线性模型
        # logging.info("============= 线性模型预测开始 ============")
        # linear_result = linear_predict_part(model_code, data)
        # results.append(linear_result)

        # 非线性模型
        logging.info("============= 非线性模型预测开始 ============")
        nonlinear_info = nonlinear_predict_part(data, model_code)
        # return {'status': 'success', 'results': {**nonlinear_info, **linear_result}}
        return {'status': 'success', 'results': nonlinear_info}
    else:
        return {"status": "error", "results": "接口请求方式异常"}


@logging_decorator
def cement_train(request):
    """
    模型训练接口
    :return:
    """
    if request.method == "GET":
        return {"status": "error", "results": "接口请求方式异常"}

    tenant_id = request.POST.get("tenantId", "1")
    start_time = request.POST.get("startTime", "2023-11-01")
    end_time = request.POST.get("endTime", "2024-11-01")
    model_code_post = request.POST.get("modelCode", None)
    product_code_post = request.POST.get("productCode", None)
    logging.info("=============================== 模型训练，租户id：{}".format(tenant_id))
    # prefix_name = settings.self_learning_collection_prefix
    # settings.self_learning_collection_name = prefix_name + str(tenant_id)
    model_type, mas_model_type = linear_train.get_model_type()  # 获取磨号、水泥品种list
    db = linear_train.get_db_handle()
    results = []
    print(" ================== {} ==================".format(mas_model_type))
    for i in range(len(model_type)):
        mill_code, product_code = model_type[i]
        if model_code_post is not None and model_code_post != mill_code[0]:
            continue
        if product_code_post is not None and product_code_post != product_code[0]:
            continue
        # try:
        logging.info("============= 开始读取自学习表 ============")
        filtered_data = linear_train.monggo_get_data(db, mill_code, product_code, None)  # 获取训练数据
        # 线性模型
        logging.info("============= 线性模型训练开始 ============")
        # linear_train_part(mas_model_type[i], filtered_data)
        # 非线性模型
        logging.info("============= 非线性模型训练开始 ============")
        model3_n1d, model3_1d, model28_n1d, model28_1d, model28_3d = nonlinear_train_part(filtered_data,
                                                                                          model_name="CatBoost",
                                                                                          mas_model_type=mas_model_type[
                                                                                              i],
                                                                                          start_time=start_time,
                                                                                          end_time=end_time)
        results.append({
            'mill_code': mill_code,
            'product_code': product_code,
            'status': 'success',
            'message': '模型训练成功'
        })
        # 循环打包上传到大模型平台上 TODO 如果上线需要将该部分进行打开
        # packaging_and_upload(mas_model_type[i])
        # delete_files_in_directory(settings.save_model_path)
        # 删除所有json文件
        # TODO 上线后该部分代码需要注释掉
        json_files = glob.glob(os.path.join(settings.save_model_path, "CatBoost/*.json"))
        for json_file in json_files:
            os.remove(json_file)
    return {'status': 'success', 'results': results}


# ==================== 异步训练模型 =====================
@logging_decorator
def async_pool_train_model(request):
    if request.method == "POST":
        return {"status": "error", "results": "接口请求方式异常"}
    model_type, mas_model_type = linear_train.get_model_type()  # 获取磨号、水泥品种list
    db = linear_train.get_db_handle()
    results = []
    print(" ================== {} ==================".format(mas_model_type))
    for i in range(len(model_type)):
        mill_code, product_code = model_type[i]
        # try:
        logging.info("============= 开始读取自学习表 ============")
        filtered_data = linear_train.monggo_get_data(db, mill_code, product_code, None)  # 获取训练数据
        # 线性模型
        logging.info("============= 线性模型训练开始 ============")
        linear_train_task = linear_train_async.delay(mas_model_type[i], filtered_data)
        # 非线性模型
        logging.info("============= 非线性模型训练开始 ============")
        nonlinear_train_task = nonlinear_train_async.delay(filtered_data,
                                                           model_name="CatBoost",
                                                           mas_model_type=mas_model_type[i])

    return {'status': 'success', "linear_model_task_id": linear_train_task.id,
            "nonlinear_model_task_id": nonlinear_train_task.id}


def delete_files_in_directory(directory):
    """
    将model文件下的所有文件进行删除
    :param directory:
    """
    # 确保目录存在
    if not os.path.exists(directory):
        print(f"目录 {directory} 不存在")
        return

    # 遍历目录中的所有文件和子目录
    for filename in os.listdir(directory):
        file_path = os.path.join(directory, filename)

        # 如果是文件，则删除
        if os.path.isfile(file_path) or os.path.islink(file_path):
            os.unlink(file_path)
            print(f"删除文件: {file_path}")
        # 如果是目录，则递归删除
        elif os.path.isdir(file_path):
            shutil.rmtree(file_path)
            print(f"删除目录: {file_path}")


# ================== API端口下需要的每个方法 ====================
# ================== 线形模型部分 ====================
def linear_predict_part(model_code: str, data: dict):
    """
    线形模型预测结果
    :param data:
    :return:
    """
    linear_result = predict_strength.cement_strength_predict(model_code, data)
    return linear_result


def linear_train_part(mas_model_type: str, filtered_data: list):
    """
    线形模型训练部分
    :param filtered_data:
    """
    linear_train.train_model(mas_model_type, filtered_data)


# ================== 非线形模型部分 =====================
def nonlinear_predict_part(data: dict, model_code: str = "CI_10001194-CI_10001073-CI_10001091"):
    """
    非线形模型预测部分
    :param data:
    :param model_code:
    :return
    """
    save_model_path = settings.save_model_path
    model_result = {}
    model_type_suffix = determine_model_use(data)
    for item in model_type_suffix:
        model_info = {
            "model_type": item,
            "model_path": os.path.join(save_model_path, "CatBoost", model_code + "_{}.pkl".format(item)),
            "first_suffix": model_type_suffix[item]['use_name'],
            "second_suffix": model_type_suffix[item]['code']
        }
        y_pred_info = predict_main(data, model_info)
        if list(y_pred_info.keys())[0] not in model_result:
            model_result.update(y_pred_info)
        else:
            first_key_name = list(y_pred_info.keys())[0]
            model_result[first_key_name].update(y_pred_info[first_key_name])
    # model_result.update(not_have_data_type)
    return {"nonlinear": model_result}


def nonlinear_train_part(filtered_data: list, model_name: str = "CatBoost", mas_model_type: str = None,
                         start_time: str = "2023-01-01", end_time: str = "2025-01-01"):
    """
    非线形模型训练部分
    :param filtered_data:
    :param model_name:
    :param mas_model_type:
    :param start_time:
    :param end_time:
    :return:
    """
    # 返回的数据类型分别为mse， r2
    model3_n1d, model3_1d, model28_n1d, model28_1d, model28_3d = nonlinear_train(filtered_data,
                                                                                 settings.save_model_path, model_name,
                                                                                 mas_model_type=mas_model_type,
                                                                                 start_time=start_time,
                                                                                 end_time=end_time)
    return model3_n1d, model3_1d, model28_n1d, model28_1d, model28_3d


# =================== 获取异步响应值 =======================
@logging_decorator
def get_index(request):
    """
    获取任务id，并解析任务结果
    """
    task_id = request.GET.get('task_id')
    async_result = AsyncResult(task_id)
    if async_result.successful():
        result = {'status': 'success', 'results': async_result.get()}
    elif async_result.failed():
        result = {'status': 'failed', 'results': "任务失败"}
    elif async_result.status == "PENDING":
        result = {'status': 'pending', 'results': "任务等待中"}
    elif async_result.status == "STARTED":
        result = {'status': 'started', 'results': "任务进行中"}
    return result


# ================================ 打包相关代码和上传到大模型 ==========================================
def rewrite_model_desc(model_path, model_code):
    """
    重写model desc文件
    :param model_path:
    :param model_code:
    :return:
    """

    # 读取 JSON 文件
    with open(model_path, 'r', encoding="utf-8") as file:
        data = json.load(file)

    data['modelCode'] = model_code
    # 保存修改后的 JSON 数据到文件
    with open(model_path, 'w', encoding="utf-8") as file:
        json.dump(data, file, indent=4, ensure_ascii=False)


def packaging_and_upload(model_code: str):
    """
    打包相关代码和上传到大模型
    :return:
    """
    # 重写json文件
    rewrite_model_desc(settings.model_desc, model_code)
    # 2. zip
    PlatformAPIHandler.pack_to_zip(directory=settings.save_zip_path, zip_file=settings.save_zip_file)
    logging.info("----zip打包完成")
    # 3. upload

    logging.info(f"----模型开始上传")
    flag, model_mini_version = PlatformAPIHandler.register_new_version_to_platform(
        directory_need_zip=settings.save_zip_path,
        zip_file_path=settings.save_zip_file,
        model_code=model_code,
        query_api=settings.query_api,
        upload_api=settings.upload_api,
        register_api=settings.register_api,
        model_name=model_code,
    )
    # 4. local
    PlatformAPIHandler.develop_in_platform(model_code=model_code,
                                           develop_type="LOCAL",
                                           model_version=model_mini_version,
                                           api_url=settings.deploy_api)

    # TODO 后期确认是否部署到生产环境
    # 5. dev
    # PlatformAPIHandler.develop_in_platform(model_code=model_code,
    #                                        develop_type="PROD",
    #                                        model_version=model_mini_version,
    #                                        api_url=settings.deploy_api)


@logging_decorator
def update_to_big_model(request):
    """

    :param request:
    :return:
    """
    if request.method == "POST":
        model_code = request.POST.get('model_code')
        packaging_and_upload(model_code=model_code)
        return {"status": "SUCCESS", "message": "{}上传成功".format(model_code)}
